"""
    this file is convert 'nii.gz', 'nii', 'mhd', 'dcm' to 'nii.gz'
"""

import logging
import os
import shutil

import SimpleITK as sitk
import nibabel as nib
import sys
import json

sys.path.append('FormatConver')


def convert_nii_to_nii_gz(nii_path, nii_gz_path):
    """
    convert nii to nii.gz
    :param nii_path: input nii path
    :param nii_gz_path: output nii.gz path
    :return:
    """
    img = nib.load(nii_path)
    nib.save(img, nii_gz_path)


def convert_mhd_to_nii_gz(mhd_path, nii_gz_path):
    """
    convert mhd to nii.gz
    :param mhd_path: input mhd path
    :param nii_gz_path: output nii.gz path
    :return:
    """
    img = sitk.ReadImage(mhd_path)
    sitk.WriteImage(img, nii_gz_path)


def convert_dcm_to_nii_gz(dcm_path, nii_gz_path):
    """
    convert dcm to nii.gz
    :param dcm_path: input dcm path
    :param nii_gz_path: output nii.gz path
    :return:
    """
    reader = sitk.ImageSeriesReader()
    dicom_names = reader.GetGDCMSeriesFileNames(dcm_path)
    reader.SetFileNames(dicom_names)
    img = reader.Execute()
    sitk.WriteImage(img, nii_gz_path)


def niigz2nii(niigz, nii):
    img = nib.load(niigz)
    nib.save(img, nii)


def convert_nii_to_mhd(niigz, mhd):
    itk_img = sitk.ReadImage(niigz)
    sitk.WriteImage(itk_img, mhd)


def convert_nii_to_dcm(niigz, dcm):
    if not os.path.isdir(dcm):
        os.mkdir(dcm)
    img = sitk.ReadImage(niigz)
    # 保存元数据
    spacing = img.GetSpacing()
    direction = img.GetDirection()
    origin = img.GetOrigin()

    with open(os.path.join(dcm, 'meta.json'), 'w') as f:
        json.dump({'spacing': spacing, 'direction': direction, 'origin': origin}, f)

    img = sitk.GetArrayFromImage(img)
    for i in range(img.shape[0]):
        select_img = sitk.GetImageFromArray(img[i])
        sitk.WriteImage(select_img, os.path.join(dcm, str(i) + '.dcm'))


def convert_dcm_to_nii(dcm, nii):
    reader = sitk.ImageSeriesReader()
    dcm_series = os.listdir(dcm)

    dcm_series = [i for i in dcm_series if i.endswith('.dcm')]

    dcm_series = sorted(dcm_series, key=lambda x: int(x.split('.')[0]))
    dcm_series = [os.path.join(dcm, i) for i in dcm_series]

    meta = json.load(open(os.path.join(dcm, 'meta.json'), 'r'))

    reader.SetFileNames(dcm_series)
    image = reader.Execute()

    image.SetSpacing(meta['spacing'])
    image.SetOrigin(meta['origin'])
    image.SetDirection(meta['direction'])

    sitk.WriteImage(image, nii)



def cheak_is_root(file_root, raw_files):
    """
    cheak is root path
    return: True or False

    Args:
        raw_files:
    """
    for file in raw_files:
        if os.path.isdir(os.path.join(file_root, file)):
            return False
    return True


def batch_convert_niigz(raw_path, out_path):
    """
        batch convert
        raw_path: input path
        out_path: output path
        this function will convert all the file in raw_path to nii.gz
    """
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    # the raw_path must is a root path
    raw_files = os.listdir(raw_path)

    if not cheak_is_root(raw_path, raw_files):
        logging.error('the raw_path must is a root path')
        return

    logging.info(f'found {len(raw_files)} files in {raw_path}')
    logging.info(f'start convert files in {raw_path} to {out_path}')

    error_files = []

    for file in raw_files:

        logging.info(f'converting {file} ...')

        file_name = file.split('.')[0]
        file_type = file.split('.')
        if len(file_type) == 2:
            file_type = file_type[1]
        else:
            file_type = file_type[-1]

        if file_type == 'gz':
            file_type = 'nii.gz'

        if file_type not in ['nii.gz', 'nii', 'mhd', 'dcm']:
            logging.error(f'{file} is not support')
            continue

        file_path = os.path.join(raw_path, file)
        if file_type == 'nii.gz':
            shutil.copy(file_path, os.path.join(out_path, file_name + '.nii.gz'))
            continue

        if file_type == 'nii':
            convert_nii_to_nii_gz(file_path, os.path.join(out_path, file_name + '.nii.gz'))
        elif file_type == 'mhd':
            convert_mhd_to_nii_gz(file_path, os.path.join(out_path, file_name + '.nii.gz'))
        elif file_type == 'dcm':
            convert_dcm_to_nii_gz(file_path, os.path.join(out_path, file_name + '.nii.gz'))
        else:
            error_files.append(file_path)
            logging.error(f'{file} is not support')

    logging.info(f'convert files in {raw_path} to {out_path} finished')
    if len(error_files) > 0:
        logging.error(f'found {len(error_files)} files is not support')
        logging.error(f'error files: {error_files}')

    return error_files


def niigz2raw(file_path, raw_file_path, out_path):
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    error_files = []

    files = os.listdir(file_path)
    files.sort()
    raw_file_path = os.listdir(raw_file_path)
    raw_file_path.sort()

    logging.info(f'found {len(files)} files in {file_path}')
    logging.info(f'start convert files in {file_path} to {out_path}')

    for f, raw in zip(files, raw_file_path):
        logging.info(f'converting {f} ...')
        file_type = raw.split('.')
        if file_type[-1] == 'gz':
            file_type = 'nii.gz'
        else:
            file_type = file_type[-1]

        if file_type == 'nii.gz':
            shutil.copy(os.path.join(file_path, f), os.path.join(out_path, raw))
        elif file_type == 'nii':
            niigz2nii(os.path.join(file_path, f), os.path.join(out_path, raw))
        elif file_type == 'mhd':
            convert_nii_to_mhd(os.path.join(file_path, f), os.path.join(out_path, raw))
        elif file_type == 'dcm':
            convert_nii_to_dcm(os.path.join(file_path, f), os.path.join(out_path, raw))
        else:
            error_files.append(os.path.join(file_path, f))
            logging.error(f'{f} is not support')

    logging.info(f'convert files in {file_path} to {out_path} finished')
    if len(error_files) > 0:
        logging.error(f'found {len(error_files)} files is not support')
        logging.error(f'error files: {error_files}')

    return error_files
